/**
 * 品牌识别应用主文件
 * 使用MobileNet预训练模型进行迁移学习，实现对Android、Apple、Windows三个品牌logo的分类
 */

// 导入必要的模块
import {getInputs} from './data.js'           // 导入数据加载函数
import * as tfvis from '@tensorflow/tfjs-vis' // 导入可视化库
import * as tf from '@tensorflow/tfjs'        // 导入TensorFlow.js
import {img2x, file2img} from './utils.js'    // 导入图像处理工具函数

// 常量定义
const NUM_CLASSES = 3                         // 分类数量（三个品牌）
const LABELS = ['android', 'apple', 'windows'] // 类别标签名称

// 页面加载完成后执行
window.onload = async() => {
    // 1. 加载训练数据
    const {inputs, labels} = await getInputs()
    
    // 2. 创建可视化界面
    const surface = tfvis.visor().surface({name: 'Brand Prediction', styles: {height: '200px'}})
    console.log(inputs, labels) // 打印数据用于调试
    
    // 3. 显示加载的图片
    inputs.forEach(img => {
        surface.drawArea.appendChild(img)
    })
    
    // 4. 加载预训练的MobileNet模型
    const MOBILENET_MODEL_PATH = 'http://127.0.0.1:8080/mobilenet/web_model/model.json'
    const mobilenet = await tf.loadLayersModel(MOBILENET_MODEL_PATH)
    mobilenet.summary() // 打印模型结构
    
    // 5. 创建截断的MobileNet模型（用于特征提取）
    // 获取倒数第二层的特征输出层
    const layer = mobilenet.getLayer('conv_pw_13_relu')
    // 创建一个新模型，输入与原MobileNet相同，输出为选定层的输出
    const truncatedMobilenet = tf.model({
        inputs: mobilenet.inputs,
        outputs: layer.output
    })
    
    // 6. 创建用于分类的自定义模型（迁移学习部分）
    const model = tf.sequential() // 创建序列模型
    
    // 添加展平层：将卷积特征从多维张量转换为一维
    model.add(tf.layers.flatten({
        inputShape: layer.outputShape.slice(1) // 移除批量维度
    }))
    
    // 添加全连接层：10个神经元，使用ReLU激活函数
    model.add(tf.layers.dense({
        units: 10,
        activation: 'relu'
    }))
    
    // 添加输出层：3个神经元对应3个类别，使用softmax激活函数输出概率分布
    model.add(tf.layers.dense({
        units: NUM_CLASSES,
        activation: 'softmax'
    }))
    
    // 7. 编译模型
    model.compile({
        loss: 'categoricalCrossentropy', // 多分类问题使用的损失函数
        optimizer: tf.train.adam(),      // Adam优化器
    })
    
    // 8. 准备训练数据
    const { xs, ys } = tf.tidy(() => {
        // 使用截断模型提取所有图片的特征
        const xs = tf.concat(inputs.map(img => truncatedMobilenet.predict(img2x(img))))
        // 将标签转换为张量
        const ys = tf.tensor(labels)
        return {xs, ys}
    })
    
    // 9. 训练模型
    await model.fit(xs, ys, {
        epochs: 20, // 训练20轮
        callbacks: tfvis.show.fitCallbacks(
            { name: '训练结果'},
            ['loss'], // 监控损失值
            {callbacks: ['onEpochEnd']} // 每轮结束时调用回调
        )
    })
    
    // 10. 定义预测函数（暴露给全局，供HTML调用）
    window.predict = async function(file) {
        // 将文件转换为图片元素
        const img = await file2img(file)
        
        // 预测过程使用tf.tidy清理中间张量
        const pred = tf.tidy(() => {
            // 预处理图片
            const x = img2x(img)
            // 使用截断模型提取特征
            const input = truncatedMobilenet.predict(x)
            // 使用自定义模型预测类别
            return model.predict(input)
        })
        
        // 获取预测结果中概率最高的类别索引
        const index = pred.argMax(1).dataSync()[0]
        
        // 使用setTimeout避免阻塞UI线程
        setTimeout(() => {
            // 显示预测结果
            alert(LABELS[index])
            // 清理预测结果张量
            pred.dispose()
        }, 0)
    }
}